import os
import cv2
import time
import argparse


def m3uUrlCheckVideo(_url):
    start_time = time.time()
    try:
        cap = cv2.VideoCapture(_url)
        ret, frame = cap.read()
        cap.release()
        return ret, time.time()-start_time
    except BaseException as _e:
        # 此处检查到异常的一律按照失败处理
        print()
        print(str(_e))
        return False, time.time()-start_time


def readFile(_path):
    # 读取出的队列
    _info_list = []
    title_tmp = None
    with open(_path, 'r', encoding='utf-8') as f:
        for line in f:
            if line.startswith("#EXTINF:-1"):
                title_tmp = line
            else:
                # find url
                if title_tmp is not None:
                    _info_list.append([title_tmp, line])
                    title_tmp = None
    return _info_list


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="m3uChecker")
    parser.add_argument("--m3u_path", default="../../source/Kodi/m3u_test.txt", type=str)
    parser.add_argument("--out_dir", default="../../source/Kodi/", type=str)
    args = parser.parse_args()

    # 获取文件名
    filename = os.path.basename(args.m3u_path)
    # 得到文件后缀
    suffix = filename.split(".")[len(filename.split(".")) - 1]
    # 拼接新的文件名
    out_filename = filename[:-(len(suffix) + 1)] + "-pick." + suffix
    out_m3u_path = os.path.join(args.out_dir, out_filename)

    # 输入文件检查
    if not os.path.exists(args.m3u_path):
        exit(-1)

    # 创建输出文件夹
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)

    info_list = readFile(args.m3u_path)

    with open(out_m3u_path, 'w', encoding='utf-8') as w:
        # 写入文件头
        w.writelines("#EXTM3U\n")
        # 链接检查
        valid_list = []
        count = 0
        for info in info_list:
            count += 1
            per = count * 100 / len(info_list)
            print("\r there is %s file, check %s file, have done %.2f%%." % (len(info_list), count, per), flush=True, end='')

            # 检查标题，重新格式化
            title = info[0][:-1]
            ret, speed = m3uUrlCheckVideo(info[1][:-1])
            if ret:
                title = title + (" [%.02fms]\n" % speed)
                valid_list.append([title, info[1]])
                w.writelines(title)
                w.writelines(info[1])
                w.flush()

        # 关闭文件
        w.close()
    print("直播源清理完毕，原有%s个，清理%s个，剩余%s个" % (len(info_list), len(info_list) - len(valid_list), len(valid_list)))